# coding=utf-8
import shutil

import tensorflow as tf
import numpy as np
import Config
from utils import metrics_utils

"""
    中文事件提取 卷积神经网络模型
    john.zhang 2016-12-15
"""


class ace_cnn_model():
    """
    中文事件的发现 由以后候选词和一个触发词是否构成事件 这里是另外一种比较简单的模型
    Relation Classification via Convolutional Deep Neural Network
    这篇文章句子特征用的是每一个词的上下文(包含当前词)作为整体的输入
    """

    def __init__(self, sentence_length=30, event_num_labels=10, role_num_labels=6, windows=3, vocab_size=2048,
                 word_embedding_size=10, event_type_embedding_size=5,
                 pos_embedding_size=5, filter_sizes=[3, 4, 5], filter_num=300, batch_size=2):
        """
        :param sentence_length: 输入句子的长度
        :param event_num_labels: 事件类别数目
        :param role_num_labels: 角色类别数目
        :param windows: 窗口的大小
        :param vocab_size: 训练集中词的数目
        :param word_embedding_size: 词嵌入维数
        :param pos_embedding_size: 位置嵌入维数
        :param trigger_vec: 触发词向量
        :param candidate_vec: 候选词向量
        :param filter_sizes: 滤波器尺寸
        :param filter_num: 滤波器数目
        """

        input_x = tf.placeholder(tf.int32, shape=[batch_size, sentence_length], name="input_x")
        self.input_x = input_x

        # 输入事件标签占位符
        # [batch_size, event_num_labels]
        input_event_y = tf.placeholder(tf.int32, shape=[batch_size, event_num_labels], name="input_event_y")
        self.input_event_y = input_event_y

        e = tf.arg_max(input_event_y, -1)
        e = tf.expand_dims(e, -1)
        type_e = tf.tile(e, [1, sentence_length])

        # 输入角色标签占位符
        # [batch_size, role_num_labels]
        input_role_y = tf.placeholder(tf.float32, shape=[batch_size, role_num_labels], name="input_role_y")
        self.input_role_y = input_role_y

        # 触发词位置向量
        input_t_pos = tf.placeholder(tf.int32, shape=[batch_size, sentence_length], name="input_t_pos")
        self.input_t_pos = input_t_pos
        # 候选词位置向量
        input_c_pos = tf.placeholder(tf.int32, shape=[batch_size, sentence_length], name="input_c_pos")
        self.input_c_pos = input_c_pos

        dropout_keep_prob = tf.placeholder(tf.float32, name="dropout_keep_prob")
        self.dropout_keep_prob = dropout_keep_prob

        # 候选词的上下文组成的特征向量
        input_c_context = tf.placeholder(tf.int32, shape=[batch_size, windows], name="input_c_context")
        self.input_c_context = input_c_context

        # 　触发词的上下文组成的特征向量
        input_t_context = tf.placeholder(tf.int32, shape=[batch_size, windows], name="input_t_context")
        self.input_t_context = input_t_context

        # 生成word_embedding 这部分操作相对比较简单 建议用cpu做Classification
        with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope("word_embedding_layer"):
            # todo CWF 词表 [vocab_size, embedding_size]
            W = tf.Variable(tf.random_normal(shape=[vocab_size, word_embedding_size], mean=0.0, stddev=0.5),
                            name="word_table")
            # 句子特征向量 [batch_size, sentence_length, word_embedding]
            input_word_vec = tf.nn.embedding_lookup(W, input_x)

            # todo PF
            #  根据 候选词位置 触发词位置 选取位置向量
            #  在(-sentence_length+1,sentence_length-1)之间一共2*(sentence_length-1)+1个数
            input_t_pos_t = input_t_pos + (sentence_length - 1)
            Tri_pos = tf.Variable(
                tf.random_normal(shape=[2 * (sentence_length - 1) + 1, pos_embedding_size], mean=0.0, stddev=0.5),
                name="tri_pos_table")
            input_t_pos_vec = tf.nn.embedding_lookup(Tri_pos, input_t_pos_t)

            input_c_pos_c = input_c_pos + (sentence_length - 1)
            Can_pos = tf.Variable(
                tf.random_normal(shape=[2 * (sentence_length - 1) + 1, pos_embedding_size], mean=0.0, stddev=0.5),
                name="candidate_pos_table")
            input_c_pos_vec = tf.nn.embedding_lookup(Can_pos, input_c_pos_c)

            # todo EF
            W_ef = tf.Variable(
                tf.random_normal(shape=[sentence_length, event_type_embedding_size], mean=0.0, stddev=0.5),
                name="EF_table")

            event_type_vec = tf.nn.embedding_lookup(W_ef, type_e)
            # print(event_type_vec.get_shape())

            # todo Sentence Feature Inpu将CWF PF EF构成一个整体的特征 作为卷积神经网络的输入
            input_sentence_vec = tf.concat([input_word_vec, input_t_pos_vec, input_c_pos_vec, event_type_vec], 2)

            # CNN支持4d输入 因此增加一维向量 用于表示输入通道数目
            intput_sentence_vec_expanded = tf.expand_dims(input_sentence_vec, -1)

            input_c_context_vec = tf.nn.embedding_lookup(W, input_c_context)
            # 触发词及其上下文 [batch_size, windows, word_embedding_size]
            input_t_context_vec = tf.nn.embedding_lookup(W, input_t_context)

        pooled_outputs = []
        # todo 卷积操作
        for i, filter_size in enumerate(filter_sizes):
            with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope(
                    'conv-maxpool-%s' % filter_size):
                # 这里的句子特征考虑的上当前词
                filter_shape = [filter_size, word_embedding_size + 2 * pos_embedding_size + event_type_embedding_size,
                                1, filter_num]
                W = tf.Variable(tf.truncated_normal(filter_shape, stddev=0.1), name="W")
                b = tf.Variable(tf.constant(0.1, shape=[filter_num]), name="b")
                # 卷积运算
                conv = tf.nn.conv2d(
                    intput_sentence_vec_expanded,
                    W,
                    strides=[1, 1, 1, 1],
                    padding="VALID",
                    name="conv")
                h = tf.nn.relu(tf.nn.bias_add(conv, b), name="relu")
                # 最大化池化 暂时不用动态池化
                pooled = tf.nn.max_pool(
                    h,
                    ksize=[1, sentence_length - filter_size + 1, 1, 1],
                    strides=[1, 1, 1, 1],
                    padding='VALID',
                    name="pool")
                # print("pooling", pooled.get_shape())
                pooled_outputs.append(pooled)
        # 使用到的所有滤波器数目(输出的通道数目)
        num_filters_total = filter_num * len(filter_sizes)
        # 多通道的数据合并
        h_pool = tf.concat(pooled_outputs, 3)
        print("h_pool", h_pool.get_shape())

        # todo 结合Sentence Level Feature 和Lexical Level Feature ，送入下一层分类器
        input_c_context_vec_flat = tf.reshape(input_c_context_vec, [-1, windows * word_embedding_size])
        # [batch_size, windows * word_embedding_size]
        input_t_context_vec_flat = tf.reshape(input_t_context_vec, [-1, windows * word_embedding_size])

        # [batch_size, num_filters_total]
        h_pool_flat = tf.reshape(h_pool, [-1, num_filters_total])
        input_sentence_features = h_pool_flat
        input_lexical_features = tf.concat([input_c_context_vec_flat, input_t_context_vec_flat], 1)

        # input_lexical_features = tf.reshape(input_word_vec, shape=(-1, sentence_length * word_embedding_size))

        # 总体的分类器 经过一层dropout 然后再送入softmax
        with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope('dropout'):
            input_sentence_features_dropout = tf.nn.dropout(input_sentence_features, dropout_keep_prob)
            input_lexical_features_dropout = tf.nn.dropout(input_lexical_features, dropout_keep_prob)
            all_input_fatures = tf.concat([input_lexical_features_dropout, input_sentence_features], 1)

        # 分类器
        with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope('softmax'):
            # 角色类别分类
            W2 = tf.Variable(tf.truncated_normal(
                [num_filters_total + 2 * windows * word_embedding_size, role_num_labels],
                stddev=0.1), name="W2")
            b2 = tf.Variable(tf.constant(0.1, shape=[role_num_labels]), name="b2")

            xw2 = tf.nn.xw_plus_b(all_input_fatures, W2, b2)
            scores_role = tf.nn.softmax(xw2, name="scores_role")
            predicts_role = tf.argmax(scores_role, axis=1, name="predicts_role")
            self.scores_role = scores_role
            self.predicts_role = predicts_role

        # 模型的代价函数 交叉熵代价函数
        with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope('loss'):
            # 角色的交叉熵代价函数
            entropy_role = tf.nn.softmax_cross_entropy_with_logits_v2(logits=scores_role, labels=input_role_y)
            loss_role = tf.reduce_mean(entropy_role)
            self.loss_role = loss_role
            tf.summary.scalar('loss', self.loss_role)
        # 准确度 用于每一次训练时调用
        with tf.device('/job:localhost/replica:0/task:0/device:CPU:0'), tf.name_scope("accuracy"):
            # 角色的准确度
            correct_role = tf.equal(predicts_role, tf.argmax(input_role_y, 1))
            accuracy_role = tf.reduce_mean(tf.cast(correct_role, "float"), name="accuracy_role")
            self.accuracy_role = accuracy_role


"""
测试程序 john.zhang 2016-11-28 已验证 程序可以运行
"""
# 数据集配置参数
from DataSets_argument import datasets
import time
import datetime
import os

# 文件信息
file = '../data/dmcnn_argument_data.txt'
# file = 'train_eval_data/datas_ace.txt'
# 生成文件的存储位置
store_path = "acdmcnn_data.txte_data_2016_12_02"
# batch_size的大小
data_batch_size = 120
# 句子的最大长度
max_sequence_length = 25
# 选取的上下文窗口的大小
windows = 3
# 数据集
datas = datasets(file=file, store_path=store_path, batch_size=data_batch_size, max_sequence_length=max_sequence_length,
                 windows=windows)
# 神经网络模型的一些参数
# 模型的最大长度
sentence_length = max_sequence_length
event_num_labels = Config.labels_event_size  # 事件类别数目
role_num_labels = Config.labels_role_size  # 角色类别数目
vocab_size = datas.words_size  # 训练集中词的数目
word_embedding_size = Config.word_embedding_size  # 词嵌入维数
pos_embedding_size = 10  # 位置嵌入维数
filter_sizes = [3, 4, 5]  # 滤波器大小
filter_num = 300  # 滤波器大小
batch_size = None  # tensorflow中支持为定义长度的标记符合
lr = 1e-2  # 学习率
num_epochs = 100
save_predicateMode_path = "./ace_cnn_model_02/model/model.ckpt"

with tf.Graph().as_default():
    sess = tf.Session()
    with sess.as_default():
        # 模型文件
        model = ace_cnn_model(sentence_length=sentence_length,
                              event_num_labels=event_num_labels,
                              role_num_labels=role_num_labels,
                              vocab_size=vocab_size,
                              word_embedding_size=word_embedding_size,
                              pos_embedding_size=pos_embedding_size,
                              filter_sizes=filter_sizes,
                              filter_num=filter_num,
                              batch_size=batch_size)

        optimizer_role = tf.train.AdamOptimizer(1e-3)
        grads_and_vars_role = optimizer_role.compute_gradients(model.loss_role)  # 角色优化器
        train_op_role = optimizer_role.apply_gradients(grads_and_vars_role)

        timestamp = str(int(time.time()))
        out_dir = os.path.abspath(os.path.join(os.path.curdir, "ace_cnn_model_02", timestamp))
        print("Writing to {}\n".format(out_dir))
        checkpoint_dir = os.path.abspath(os.path.join(out_dir, "checkpoints"))
        checkpoint_prefix = os.path.join(checkpoint_dir, "model")
        if not os.path.exists(checkpoint_dir):
            os.makedirs(checkpoint_dir)
        saver = tf.train.Saver(tf.global_variables(), max_to_keep=100)  # 最大支持存储100个模型
        # 初始化　变量
        sess.run(tf.initialize_all_variables())
        # 设置summary
        merged = tf.summary.merge_all()
        if os.path.exists('../log/argument'):
            shutil.rmtree('../log/argument')
        train_writer = tf.summary.FileWriter('../log/argument', sess.graph)


        def train_step(input_x, input_event_y, input_role_y, input_t, input_c, input_t_pos, input_c_pos,
                       dropout_keep_prob
                       , sentence_features, input_t_context, input_c_context, train_op, loss, accuracy, predicts, epoch,
                       step, stype):
            feed_dict = {
                model.input_x: input_x,  # 句子级别特征
                model.input_event_y: input_event_y,  # 事件类别
                model.input_role_y: input_role_y,  # 角色类别
                model.input_t_pos: input_t_pos,  # 触发词位置向量
                model.input_c_pos: input_c_pos,  # 候选词位置向量
                model.dropout_keep_prob: dropout_keep_prob,  # drop_out 概率
                model.input_t_context: input_t_context,  # 触发词以及其上下文
                model.input_c_context: input_c_context  # 候选词以及上下文
            }
            _, loss, accuracy, predicts,train_summary = sess.run(
                [train_op, loss, accuracy, predicts,merged],
                feed_dict)
            time_str = datetime.datetime.now().isoformat()
            train_writer.add_summary(train_summary, step)
            print("stype:{}, epoch {}: , loss {:g}, acc {:g}".format(stype, epoch, loss, accuracy))
            if step % 10 == 0:
                y_true = [np.argmax(role_y) for role_y in input_role_y]
                metrics_utils.role_evaluate(y_true=y_true, y_pred=predicts)


        # 测试阶段不需要计算梯度 也不需要进行权值更新 仅仅需要计算acc的值
        def eval_step(input_x, input_event_y, input_role_y, input_t, input_c, input_t_pos, input_c_pos,
                      dropout_keep_prob
                      , sentence_features, input_t_context, input_c_context, accuracy, predicts, stype):
            feed_dict = {
                model.input_x: input_x,  # 句子级别特征
                model.input_event_y: input_event_y,  # 事件类别
                model.input_role_y: input_role_y,  # 角色类别
                model.input_t_pos: input_t_pos,  # 触发词位置向量
                model.input_c_pos: input_c_pos,  # 候选词位置向量
                model.dropout_keep_prob: dropout_keep_prob,  # drop_out 概率
                model.input_t_context: input_t_context,  # 触发词以及其上下文
                model.input_c_context: input_c_context  # 候选词以及上下文
            }
            accuracy, predicts = sess.run([accuracy, predicts], feed_dict)
            print("{} eval accuracy:{}".format(stype, accuracy))
            y_true = [np.argmax(role_y) for role_y in input_role_y]
            metrics_utils.role_evaluate(y_true=y_true, y_pred=predicts)
            return predicts

        totol_step = 0
        for i in range(num_epochs):
            for j in range(datas.instances_size // data_batch_size):
                x, t, c, y_e, y_r, pos_c, pos_t, sentences_f, c_context, t_context, _, _, _ = datas.next_cnn_data()
                #     角色类型预测
                totol_step += 1
                train_step(input_x=x, input_event_y=y_e, input_role_y=y_r,
                           input_t=t, input_c=c, input_c_pos=pos_c, input_t_pos=pos_t,
                           dropout_keep_prob=0.5, sentence_features=sentences_f,
                           input_t_context=t_context, input_c_context=c_context,
                           train_op=train_op_role, loss=model.loss_role, accuracy=model.accuracy_role,
                           predicts=model.predicts_role,
                           epoch=i, step=totol_step, stype="role"
                           )
        saver.save(sess, save_predicateMode_path)
        saver.restore(sess, save_predicateMode_path)
        print \
            ("----------------------------华丽的分割线-----------------------------------------")
        x, t, c, y_e, y_r, pos_c, pos_t, sentences_f, c_context, t_context, _1, _2, _3 = datas.eval_cnn_data()
        # 角色类型预测

        predicts_role = eval_step(input_x=x, input_event_y=y_e, input_role_y=y_r, input_t=t, input_c=c,
                                  input_c_pos=pos_c, input_t_pos=pos_t,
                                  dropout_keep_prob=1.0, sentence_features=sentences_f, input_t_context=t_context,
                                  input_c_context=c_context, accuracy=model.accuracy_role, predicts=model.predicts_role,
                                  stype="role")
        convert_event = {0: "非事件", 1: "statement", 2: "emergency", 3: "perception", 4: "stateChange",
                         5: "operation", 6: "action", 7: "movement"
                         }  # 事件类型准换
        convert_role = {0: "其它角色", 1: "施事角色", 2: "受事角色", 3: "时间角色", 4: "地点角色"}  # 角色类型转换
        # 输出测试结果
        for i in range(len(x)):
            print("输入数据：{}".format(", ".join(map(lambda h: datas.all_words[h], x[i]))))
            print("触发词：{}".format(", ".join(map(lambda h: datas.all_words[h], t[i]))))
            print("候选词：{}".format(", ".join(map(lambda h: datas.all_words[h], c[i]))))
            print("预测角色类别:{}".format(convert_role[predicts_role[i]]))
            print("实际角色类别:{}".format(convert_role[np.argmax(y_r[i])]))
